import matplotlib as mpl
import matplotlib.pyplot as plt

import numpy as np
import sklearn  # 注意
import pandas as pd
import os
import sys
import time
import tensorflow as tf

from tensorflow import keras

from tensorflow.keras import Sequential, layers, datasets


# 加载手写数据集文件
def preprocess(x, y):
    """
    预处理函数
    """
    # [b, 28, 28], [b]
    x = tf.cast(x, dtype=tf.float32) / 255.
    y = tf.cast(y, dtype=tf.int32)
    y = tf.one_hot(y, depth=10)
    return x, y


(x_train, y_train), (x_test, y_test) = datasets.mnist.load_data()  # 加载手写数据集数据
x_train, x_test = x_train / 255.0, x_test / 255.0
# 增加一个维度
x_train = np.expand_dims(x_train, axis=3)
x_test = np.expand_dims(x_test, axis=3)

# BUILD CONVOLUTIONAL NEURAL NETWORKS
nets = 3
model = [0] * nets

for j in range(3):
    model[j] = Sequential()
    model[j].add(keras.layers.Conv2D(24, kernel_size=5, padding='same', activation='relu',
                                     input_shape=(28, 28, 1)))
    model[j].add(keras.layers.MaxPool2D())
    if j > 0:
        model[j].add(keras.layers.Conv2D(48, kernel_size=5, padding='same', activation='relu'))
        model[j].add(keras.layers.MaxPool2D())
    if j > 1:
        model[j].add(keras.layers.Conv2D(64, kernel_size=5, padding='same', activation='relu'))
        model[j].add(keras.layers.MaxPool2D(padding='same'))
    model[j].add(keras.layers.Flatten())
    model[j].add(keras.layers.Dense(256, activation='relu'))
    model[j].add(keras.layers.Dense(10, activation='softmax'))
    model[j].compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])

# 训练模型
epoch = 5
history = [0] * nets

for j in range(nets):

    history[j] = model[j].fit(x=x_test, y=y_test, batch_size=256, epochs=epoch, validation_data=(x_test, y_test),
                              verbose=0)
    #print("CNN {0}: Epochs={1:d}, Train accuracy={2:.5f}, Validation accuracy={3:.5f}". \
    #      format(j, epoch, history.history['sparse_categorical_accuracy'],
     #            history.history['val_sparse_categorical_accuracy']))
# model[2].summary()
